/**
 * args parser for qualiti tools
 * */
#ifndef __ARGS_PARSER
#define __ARGS_PARSER

#include <string>
#include <vector>
#include <iostream>
#ifdef _MSC_VER
#include "../common/windows/getopt.h"
#else
#include <getopt.h>
#endif
namespace calibrate_tool{

typedef struct _Args Args;

/**
 * \brief Struct to maintain comman-line arguments.
 * */
struct  _Args
{
    int batch{1};
    int int8{-1};
    int input_h{384};
    int input_w{640};
    std::string onnxFile{"../appPlan1/app1_2/models/Det_FullPic_PersonCar_384x640_zyf.onnx"};
    std::string outFile{"out_int8.engine"};
    std::string dataDir{"calibration_data/PersonCarCalibrationData/"};
    std::string calib_table_name{"int8_calib.table"};
    std::string input_blob_name{"images"};
};

bool parseArgs(Args& args, int argc, char* argv[]){
    while (1)
    {
        int arg;
        static struct option long_options[] = {
            {"batch", required_argument, 0, 'b'}, {"useint8", required_argument, 0, 'u'}, 
            {"input_h", required_argument, 0, 'h'}, {"input_w", required_argument, 0, 'w'},
            {"onnxfile", required_argument, 0, 'i'}, {"outfile", required_argument, 0, 'o'},
            {"datadir", required_argument, 0, 'd'}, {"calib_table", required_argument, 0, 'c'},
            {"input_blob", required_argument, 0, 'n'}, {nullptr, 0, nullptr, 0}};
        int option_index = 0;
        arg = getopt_long(argc, argv, "b:u:h:w:i:o:d:c:n:", long_options, NULL);
        if (arg == -1)
        {
            break;
        }

        switch (arg)
        {
        case 'b':
            if(optarg)
            {
                args.batch = std::stoi(optarg);
            }
            break;
        case 'u': args.int8 = std::stoi(optarg); break;
        case 'h': args.input_h = std::stoi(optarg); break;
        case 'w': args.input_w = std::stoi(optarg); break;
        case 'i': args.onnxFile = optarg; break;
        case 'o': args.outFile = optarg;break;
        case 'd':
            if (optarg)
            {
                args.dataDir = optarg;
            }
            else
            {
                std::cerr << "ERROR: --datadir requires option argument" << std::endl;
                return false;
            }
            break;
        case 'c': args.calib_table_name = optarg; break;
        case 'n': args.input_blob_name = optarg; break;
        default: return false;
        }
    }
    return true;
};

}; // namespace calibrate_tool

namespace infer_utils{

typedef struct _Args Args;
struct _Args{
    int batch{1};
    std::string datadir;
    int input_h{224};
    int input_w{224};
    std::string input_blob_name;
    std::string output_blob_name;

    std::string engine1;
    std::string engine2;
};

void printUsage(){
    std::cout << "Usage: \n"
              << "[--batch ] : models batchsize default 1. \n"
              << "[--engine1 ] : the first engine's pos\n"
              << "[--engine2 ] : the second engine's pos\n"
              << "[--datadir ] : the imgs's dir engine will infer \n"
              << "[--input_h ] : the engine's input size h\n"
              << "[--input_w ] : the engine's input size w\n"
              << "[--input_blob_name ] : the engine's input_blob_name \n"
              << "[--output_blob_name ] : the engine's output_blob_name \n";
}

void printArgs(Args& args){
    std::cout << "=======args==========\n"
              << "batch   : " <<args.batch << "\n"
              << "engine1 : " << args.engine1 << "\n"
              << "engine2 : " << args.engine2 << "\n"
              << "datadir : " << args.datadir << "\n"
              << "input_h : " << args.input_h << "\n"
              << "input_w : " << args.input_w << "\n"
              << "input_blob_name  : " << args.input_blob_name  << "\n"
              << "output_blob_name : " << args.output_blob_name << "\n";
}

bool parseArgs(Args& args, int argc, char* argv[]){
    while (1)
    {
        int arg;
        static struct option long_options[] = {
            {"batch", required_argument, 0, 'b'}, {"datadir", required_argument, 0, 'd'},
            {"input_h", required_argument, 0, 'h'}, {"input_w", required_argument, 0, 'w'},
            {"engine1", required_argument, 0, 'i'}, {"engine2", required_argument, 0, 'o'},
            {"input_blob_name", required_argument, 0, 'n'}, {"output_blob_name", required_argument, 0, 'm'}, 

            {nullptr, 0, nullptr, 0}};
        int option_index = 0;
        arg = getopt_long(argc, argv, "b:d:h:w:i:o:", long_options, NULL);
        if (arg == -1)
        {
            break;
        }

        switch (arg)
        {
        case 'b':
            if(optarg)
            {
                args.batch = std::stoi(optarg);
            }
            break;
        case 'h': args.input_h = std::stoi(optarg); break;
        case 'w': args.input_w = std::stoi(optarg); break;
        case 'i': args.engine1 = optarg; break;
        case 'o': args.engine2 = optarg; break;
        case 'n': args.input_blob_name  = optarg; break;
        case 'm': args.output_blob_name = optarg; break;
        case 'd':
            if (optarg)
            {
                args.datadir = optarg;
            }
            else
            {
                std::cerr << "ERROR: --datadir requires option argument" << std::endl;
                return false;
            }
            break;
        default: return false;
        }
    }
    return true;
}

};

#endif